import os
import torch
import numpy as np
from torch.utils.data import DataLoader
from torchvision import transforms
from PIL import Image
import argparse
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
from dataset import CatDogDataset
from net import CatDogClassifier


def evaluate_model(model, test_loader, device):
    model.eval()
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        for images, labels in test_loader:
            images = images.to(device)
            labels = labels.to(device)
            
            outputs = model(images)
            _, preds = torch.max(outputs, 1)
            
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())
    
    return all_labels, all_preds

def plot_confusion_matrix(cm, classes, save_path='confusion_matrix.png'):
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=classes, yticklabels=classes)
    plt.xlabel('Predicted')
    plt.ylabel('Actual')
    plt.title('Confusion Matrix')
    plt.savefig(save_path)
    plt.close()

def main():
    parser = argparse.ArgumentParser(description='Test the CatDog Classifier')
    parser.add_argument('--test-dir', type=str, default=r"data\test", help='Directory of test data')
    parser.add_argument('--model-path', type=str,default=r"saved_models\fold4_epoch0_acc0.9525.pth", help='Path to the trained model')
    parser.add_argument('--batch-size', type=int, default=32, help='Batch size for testing')
    parser.add_argument('--save-cm', action='store_true', default=True, help='Save confusion matrix plot')
    args = parser.parse_args()

    # 设备设置
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    # 数据预处理 (与训练时相同)
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    # 加载测试集
    test_dataset = CatDogDataset(args.test_dir, transform=transform)
    test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False)

    # 加载模型
    model = CatDogClassifier().to(device)
    model.load_state_dict(torch.load(args.model_path))
    print(f"Loaded model from: {args.model_path}")

    # 评估模型
    true_labels, pred_labels = evaluate_model(model, test_loader, device)

    # 计算各项指标
    print("\nClassification Report:")
    print(classification_report(true_labels, pred_labels, target_names=['cat', 'dog'], digits=4))

    # 计算混淆矩阵
    cm = confusion_matrix(true_labels, pred_labels)
    print("\nConfusion Matrix:")
    print(cm)

    # 可视化混淆矩阵
    if args.save_cm:
        plot_confusion_matrix(cm, ['cat', 'dog'])
        print("Confusion matrix plot saved as 'confusion_matrix.png'")

    # 计算并打印各项指标
    tn, fp, fn, tp = cm.ravel()
    
    accuracy = (tp + tn) / (tp + tn + fp + fn)
    precision = tp / (tp + fp) if (tp + fp) > 0 else 0
    recall = tp / (tp + fn) if (tp + fn) > 0 else 0
    f1 = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0
    
    print("\nDetailed Metrics:")
    print(f"Accuracy:  {accuracy:.4f}")
    print(f"Precision: {precision:.4f}")
    print(f"Recall:    {recall:.4f}")
    print(f"F1 Score:  {f1:.4f}")

if __name__ == '__main__':
    main()